#pragma once
#include <poll.h>
#include <list>
#include <unordered_map>
#include "Common.hpp"
#include "Socket.hpp"
using namespace SocketModule;

class PollServer : public NoCopy
{
private:
    void ListenSocketHandler()
    {
        std::shared_ptr<TCPConnectSocket> connect_socket = _listen_socket->Accept();
        if (!connect_socket)
        { // 检查Accept是否成功
            LOG(LogLevel::ERROR) << "Accept failed";
            return;
        }

        int connect_sockfd = connect_socket->SockFd();
        RegisterFd(connect_sockfd, POLLIN);
        _connect_socket[connect_sockfd] = connect_socket;
    }

    void ConnectSocketHandler(int fd)
    {
        auto it = _connect_socket.find(fd);
        if (it == _connect_socket.end())
        {
            LOG(LogLevel::ERROR) << "Invalid socket fd: " << fd;
            return;
        }

        std::string message;
        int n = it->second->Receive(message);
        std::string client = it->second->Addr().Info();

        if (n > 0)
        {
            // 正常收到消息
            std::cout << "Client[" << client << "] say# " << message << std::endl;
        }
        else if (n == 0)
        {
            // 客户端断开连接
            LOG(LogLevel::INFO) << "Client[" << client << "]已断开连接...";
            _connect_socket.erase(it);
            DeleteFd(fd);
        }
        else
        {
            // 出错
            LOG(LogLevel::ERROR) << "Receive: 接收Client[" << client << "]的数据失败! ";
            _connect_socket.erase(it);
            DeleteFd(fd);
        }
    }

    void Dispatch()
    {
        for (auto &poll : _fds)
        {
            // 跳过无效的fd或没有事件的fd
            if (poll.fd == _default_socket_fd || !(poll.revents & POLLIN))
                continue;

            if (poll.fd == _listen_socket->SockFd())
            {
                // 监听套接字
                ListenSocketHandler();
            }
            else
            {
                // 连接套接字
                ConnectSocketHandler(poll.fd);
            }
        }
    }

    void RegisterFd(int fd, short events)
    {
        // 尝试找到一个空闲位置
        for (auto &poll : _fds)
        {
            if (poll.fd == _default_socket_fd)
            {
                poll.fd = fd;
                poll.events = events;
                return;
            }
        }

        // 如果没有空闲位置，则添加新元素
        _fds.push_back((pollfd){fd, events, 0});
    }

    void DeleteFd(int fd)
    {
        for (auto &poll : _fds)
        {
            if (poll.fd == fd)
            {
                poll.fd = _default_socket_fd;
                poll.events = 0;
                poll.revents = 0;
                break;
            }
        }
    }

public:
    PollServer(in_port_t port)
        : _isrunning(false), _listen_socket(std::make_shared<TCPListenSocket>(port))
    {
        if (!_listen_socket || _listen_socket->SockFd() < 0)
        {
            LOG(LogLevel::FATAL) << "PollServer: 初始化监听套接字失败! ";
            exit(EXIT_FAILURE);
        }

        int listen_sockfd = _listen_socket->SockFd();
        RegisterFd(listen_sockfd, POLLIN);
    }

    void Run()
    {
        _isrunning = true;
        while (_isrunning)
        {
            // 使用实际有效的fd数量
            int n = poll(_fds.data(), _fds.size(), 10000);

            if (n < 0)
            {
                // 处理错误，EINTR是可恢复的
                if (errno == EINTR)
                    continue;
                LOG(LogLevel::FATAL) << "poll error: " << strerror(errno);
                break;
            }
            else if (n == 0)
            {
                // 超时
                LOG(LogLevel::DEBUG) << "poll timeout...";
            }
            else
            {
                // 有事件就绪
                LOG(LogLevel::INFO) << "有" << n << "个事件就绪, 即将处理...";
                Dispatch();
            }
        }
        _isrunning = false;
    }

    ~PollServer()
    {
        _isrunning = false;
    }

private:
    bool _isrunning;
    std::shared_ptr<TCPListenSocket> _listen_socket;                            // 监听套接字
    std::unordered_map<int, std::shared_ptr<TCPConnectSocket>> _connect_socket; // 需要poll等待的连接套接字
    std::vector<pollfd> _fds;                                                   // poll等待列表
    static const int _default_socket_fd;
};

const int PollServer::_default_socket_fd = -1;